package leetcode_day._2022._202201._1120;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;

/**
 * @author yzh
 * @version 1.0
 * @date 2022/1/14 9:53
 * 查找和最小的 k 对数字
 * 算法：数组、多路并归
 * 使用数组计算 nums1 在 nums2 走的步数
 */
public class _14_373 {

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int n = nums1.length, m = nums2.length;
        if (n == 0 || m == 0) return new ArrayList<>(0);
        if (k > n * m) k = n * m;
        List<List<Integer>> ans = new ArrayList<>(k);
        int[] steps = new int[n];
        Arrays.fill(steps, 0);
        for (int i = 0; i < k; i++) {
            int min = Integer.MAX_VALUE, minStepIdx = 0;
            for (int j = 0; j < n; j++) {
                if (steps[j] < m && nums1[j] + nums2[steps[j]] < min) {
                    min = nums1[j] + nums2[steps[j]];
                    minStepIdx = j;
                }
            }
            ans.add(Arrays.asList(nums1[minStepIdx], nums2[steps[minStepIdx]]));
            steps[minStepIdx]++;
        }
        return ans;
    }

    public List<List<Integer>> kSmallestPairs_(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<>(k, (a, b) -> nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]]);
        int m = nums1.length, n = nums2.length;
        for (int i = 0; i < Math.min(m, k); i++) queue.offer(new int[]{i, 0});
        List<List<Integer>> ans = new ArrayList<>();
        while (k-- > 0 && !queue.isEmpty()) {
            final int[] poll = queue.poll();
            ans.add(Arrays.asList(nums1[poll[0]], nums2[poll[1]]));
            if (poll[1] + 1 < n) queue.offer(new int[]{poll[0], poll[1] + 1});
        }
        return ans;
    }

}
